import io
import draw
from flask import Flask, request, send_from_directory
app = Flask(__name__)


@app.route('/')
def index():
    return send_from_directory('templates', "index.html")


@app.route('/whole', methods=['POST'])
def whole():
    data = request.get_json()
    algorithm = data['algorithm']  # 算法
    dataset = data['dataset']  # 数据集
    data_type = data['data_type']  # 训练集or测试集

    # 调取画图接口
    canvas = draw.whole_picture(algorithm, dataset, data_type)
    buffer = io.BytesIO()
    canvas.print_png(buffer)
    data = buffer.getvalue()
    buffer.close()
    res = app.make_response(data)
    res.headers["Content-Type"] = "image/png"
    return res


@app.route('/single', methods=['POST'])
def single():
    data = request.get_json()
    algorithm = data['algorithm']  # 算法
    dataset = data['dataset']  # 数据集
    node = data['node']  # 攻击节点
    # 调取画图接口
    canvas = draw.single_pictue(algorithm, dataset, int(node))
    buffer = io.BytesIO()
    canvas.print_png(buffer)
    data = buffer.getvalue()
    buffer.close()
    res = app.make_response(data)
    res.headers["Content-Type"] = "image/png"
    return res


if __name__ == '__main__':
    app.run(debug=False)
